import json
import argparse
import os
import numpy as np
from tqdm import tqdm  # 引入 tqdm 库

'''
# @Time : 2024/11/23 17:46
# @Author : Shiyu He
# @University : Xinjiang University
'''


# 设置命令行参数
parser = argparse.ArgumentParser(description="Extract data from JSON file and save to separate files.")
parser.add_argument("input_json", type=str, help="Path to the input JSON file.")
parser.add_argument("output_path", type=str, help="Path to save extracted data.")
parser.add_argument("extract_length", type=int, help="Extract Count.")
args = parser.parse_args()

# 文件路径
input_json_file = args.input_json
output_path = args.output_path
extract_length = args.extract_length

# 从 JSON 文件读取数据
with open(input_json_file, "r") as file:
    data = json.load(file)

# 初始化提取结果
plane_positions = {}
plane_speeds = {}
enemy_positions = {}
enemy_speeds = {}
outputs = {}

# 解析数据
# 使用 tqdm 显示进度条，限制为最多 3000 条
for idx, entry in tqdm(enumerate(data), total=min(len(data), extract_length), desc="Extracting Data"):
    if idx >= extract_length:
        break  # 超过最大条目数时退出循环
    
    instruction = entry["instruction"]
    output = entry["output"]
    
    # 提取飞机位置和速度
    plane_pos_start = instruction.find("position is (") + len("position is (")
    plane_pos_end = instruction.find(")", plane_pos_start)
    plane_position = tuple(map(float, instruction[plane_pos_start:plane_pos_end].split(", ")))

    speed_start = instruction.find("speed is (") + len("the speed is (")
    speed_end = instruction.find(")", speed_start)
    plane_speed = tuple(map(float, instruction[speed_start:speed_end].split(", ")))
    
    # 提取敌方相对位置和速度
    enemy_pos_start = instruction.find("relative position of the enemy is (") + len("relative position of the enemy is (")
    enemy_pos_end = instruction.find(")", enemy_pos_start)
    enemy_position_relative = tuple(map(float, instruction[enemy_pos_start:enemy_pos_end].split(", ")))
    
    enemy_speed_start = instruction.find("relative speed of the enemy is (") + len("relative speed of the enemy is (")
    enemy_speed_end = instruction.find(")", enemy_speed_start)
    enemy_speed = tuple(map(float, instruction[enemy_speed_start:enemy_speed_end].split(", ")))

    # 计算敌方的绝对位置（可选）
    enemy_position = tuple(np.array(plane_position) + np.array(enemy_position_relative))
    
    # 提取控制指令 (output)
    output_start = output.find("[") + 1
    output_end = output.find("]")
    control_instruction = list(map(float, output[output_start:output_end].split(", ")))
    
    # 保存数据，使用索引作为键
    plane_positions[str(idx)] = plane_position
    plane_speeds[str(idx)] = plane_speed
    enemy_positions[str(idx)] = enemy_position
    enemy_speeds[str(idx)] = enemy_speed
    outputs[str(idx)] = control_instruction
    # 获取提取到的条数
    state_count = len(plane_positions)
    control_count = len(outputs)
# 检查条数是否一致
assert state_count == control_count, f"Length mismatch: state_count={state_count}, control_count={control_count}. Please check the input data."
# 输出提取的条数
print(f"提取到的数据条数: {state_count}")

# 保存提取的 input 和 output 数据
with open(os.path.join(output_path, "state.json"), "w") as state_file:
    json.dump({
        "plane_positions": plane_positions,
        "plane_speeds": plane_speeds,
        "enemy_positions": enemy_positions,
        "enemy_speeds": enemy_speeds
    }, state_file, indent=4)

with open(os.path.join(output_path, "control.json"), "w") as control_file:
    json.dump({"control": outputs}, control_file, indent=4)

print(f"提取完成！\n数据已经保存到: {output_path}")
